"""

This script takes in files that represent classifications of user-years into different clusters, and aggregates their data into:

	1. Count of users in each cluster by year, e.g.
		year	c1	c2	c3
		2003	100	50	50
		2004	200	100	100
		2005	250	150	150
		...
		2012	300	200	200
	
	2. Proportion of users in each cluster by year, e.g.
		year	p1	p2	p3
		2003	0.9	0.05	0.05
		2004	0.8	0.1	0.1
		...
		2012	0.5	0.25	0.25	

Each block of statistics gets its own output file.

Specify an input directory to read input files from.
Specify an output directory to write output files to.
Specify the number of clusters in the model. This probably isn't necessary but is a sanity check in case no users are classified to a cluster.

There are lots of parameters for which elements to look for in this script that are not passed as arguments. The defaults here should be good for 
our project.

This script is well commented enough that reading through the code should answer any questions about how these statistics are compiled.

"""

import sys
import os
import datetime

def main():

	inputDir = sys.argv[1]
	outputDir = sys.argv[2]
	numclusters = int(sys.argv[3])


	# Make a list of input files and open them.

	inputFileList = os.listdir(inputDir)
	inputFileHandles = []
	
	for inputFileName in inputFileList:
		
		inputFilePath = inputDir+'/'+inputFileName
		
		if os.path.isfile(inputFilePath):
			inputFileHandles.append(open(inputFilePath, 'r'))

	# Initialize count and proportion by year

	year = {}			# Address year[year][cluster]
	proportion_year = {}		# Address proportion_year[year][cluster]

	start_year = 2003		# The first year we look at in the data, inclusive.
	end_year = 2012			# The last year we look at in the data, inclusive.

	for year_index in range(start_year, end_year + 1):
	
		year[year_index] = {}
		proportion_year[year_index] = {}

		for cluster_index in range (1, numclusters + 1):

			year[year_index][cluster_index] = 0
			proportion_year[year_index][cluster_index] = None


	# Compile counts

	print timestamp() + " Processing input files in " + inputDir

	total_entries = 0		# The total number of entries we look at

	for inputFileHandle in inputFileHandles:

		print timestamp() + " Processing input file " + inputFileHandle.name

		for line in inputFileHandle:

			fields = parseLine(line)

			# Update count by year

			year[fields['year']][fields['cluster']] += 1

			
			# Increment the total entry counter

			total_entries += 1


		print timestamp() + " Finished processing input file " + inputFileHandle.name


	print timestamp() + " Finished processing input files in " + inputDir


	# Compute proportions

	print timestamp() + " Compiling proportions"


	# Compute year proportion

	year_total = 0		# Sanity check to make sure we counted all entries

	for year_index in year:

		single_year_total = 0		# The total for one year that we compute the proportion with

		# Figure out the total users in this year

		for cluster_index in year[year_index]:

			year_total += year[year_index][cluster_index]
			single_year_total += year[year_index][cluster_index]

		# Compute the proportion based on the total users in this year

		for cluster_index in year[year_index]:

			if single_year_total != 0:
				proportion_year[year_index][cluster_index] = float(year[year_index][cluster_index]) / float(single_year_total)
			else:
				proportion_year[year_index][cluster_index] = 0

	assert(year_total == total_entries), "Year count is missing " + str(total_entries - year_total) + " entries"

	
	# Write year count statistics file
	
	year_filename = 'year.stat'						# Filename to write year count statistics to
	year_filepath = outputDir + '/' +year_filename

	print timestamp() + " Writing year count statistics file to " + year_filepath

	year_filehandle = open(year_filepath, 'w')
	
	for year_index in sorted(year):

		output = [str(year_index)]

		for cluster_index in sorted(year[year_index]):

			output.append(str(year[year_index][cluster_index]))

		year_filehandle.write('\t'.join(output) + '\n')

	year_filehandle.close()

	print timestamp() + " Finished writing year count statistics file to " + year_filepath

	
	# Write year proportion statistics file

	proportion_year_filename = 'year_proportion.stat'			# Filename to write year proportion statistics to
	proportion_year_filepath = outputDir + '/' + proportion_year_filename

	print timestamp() + " Writing year proportion statistics file to " + proportion_year_filepath

	proportion_year_filehandle = open(proportion_year_filepath, 'w')

	for year_index in sorted(proportion_year):

		output = [str(year_index)]

		for cluster_index in sorted(proportion_year[year_index]):

			output.append(str(proportion_year[year_index][cluster_index]))

		proportion_year_filehandle.write('\t'.join(output) + '\n')

	proportion_year_filehandle.close()

	print timestamp() + " Finished writing year proportion statistics file to " + proportion_year_filepath


def parseLine(line):
	
	tokens = line.strip().split('\t')

	fields = {}
	fields['user'] = tokens[0]
	fields['offset'] = int(tokens[1])
	fields['year'] = int(tokens[2])
	fields['cluster'] = int(tokens[3])

	return fields

def timestamp():
	
	return datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')

main()
